package com.example.demo.interceptor;

import com.example.demo.annotation.RateLimiter;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;

@Configuration
public class CustomInterceptor implements HandlerInterceptor {

    @Resource
    private RedissonClient client;

    private final ConcurrentHashMap<String,RRateLimiter> map = new ConcurrentHashMap<>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod){
            Method method = ((HandlerMethod) handler).getMethod();

            Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
            String userName = authentication.getName();

            String path=request.getServletPath();
            if (method.isAnnotationPresent(RateLimiter.class)){
                RateLimiter rateLimiter = method.getAnnotation(RateLimiter.class);

                RRateLimiter limiter;
                if (userName!=null && !userName.equals("anonymousUser")){
                    int userLimit = rateLimiter.userLimit();

                    String key=userName+path+"user-rate-limiter";
                    if (!map.containsKey(key)){
                        limiter = client.getRateLimiter(key);
                        limiter.setRate(RateType.PER_CLIENT,userLimit,1, RateIntervalUnit.MINUTES);

                        map.put(key,limiter);
                    }else {
                        limiter = map.get(key);
                    }

                }else {
                    int ipLimiter = rateLimiter.ipLimit();

                    String remoteAddr = request.getRemoteAddr();
                    String key=remoteAddr+path+"ip-rate-limiter";
                    if (!map.containsKey(key)){
                        limiter = client.getRateLimiter(key);
                        limiter.setRate(RateType.PER_CLIENT,ipLimiter,1, RateIntervalUnit.MINUTES);

                        map.put(key,limiter);
                    }else {
                        limiter = map.get(key);
                    }
                }

                if (limiter.tryAcquire()){
                    return true;
                }else {
                    throw new RuntimeException("请稍后重试");
                }
            }
        }

        return true;
    }
}
